use futures::stream::FuturesUnordered;
use futures::StreamExt;
use hashbrown::HashMap;
use xactor::{message, Actor, Addr, Context, Handler, Message, Result, Service};

#[async_trait::async_trait]
pub trait ShutDownAddr: Send + Sync {
    async fn shutdown(&self);
}

#[async_trait::async_trait]
impl<T: Actor> ShutDownAddr for Addr<T> {
    async fn shutdown(&self) {
        let _ = self.stop(None);
        self.wait_for_stop().await;
    }
}

#[derive(Default)]
pub struct ShutdownActor {
    actors: HashMap<u64, Box<dyn ShutDownAddr>>,
}

impl Service for ShutdownActor {}

#[async_trait::async_trait]
impl Actor for ShutdownActor {
    fn name(&self) -> &'static str {
        "ShutdownActor"
    }
}

pub struct Register<T: Actor>(pub Addr<T>);

impl<T: Actor> Message for Register<T> {
    type Result = ();
}

#[async_trait::async_trait]
impl<T: Actor> Handler<Register<T>> for ShutdownActor {
    async fn handle(&mut self, _ctx: &mut Context<Self>, msg: Register<T>) -> Result<()> {
        let actor_id = msg.0.actor_id();
        self.actors.insert(actor_id, Box::new(msg.0));
        debug!("{} register {}", self.name(), actor_id);
        Ok(())
    }
}

#[message]
pub struct UnRegister(pub u64);

#[async_trait::async_trait]
impl Handler<UnRegister> for ShutdownActor {
    async fn handle(&mut self, _ctx: &mut Context<Self>, msg: UnRegister) -> Result<()> {
        self.actors.remove(&msg.0);
        debug!("{} unregister {}", self.name(), msg.0);
        Ok(())
    }
}

#[message]
pub struct ShutDown;
#[async_trait::async_trait]
impl Handler<ShutDown> for ShutdownActor {
    async fn handle(&mut self, ctx: &mut Context<Self>, _: ShutDown) -> Result<()> {
        let mut shutdown_stream: FuturesUnordered<_> = self
            .actors
            .iter()
            .map(|(_, actor)| actor.shutdown())
            .collect();

        while shutdown_stream.next().await.is_some() {}

        drop(shutdown_stream);

        self.actors.clear();

        ctx.stop(None);
        Ok(())
    }
}
